/*
 * Copyright (c) 2020-2099 the original author or authors. All rights reserve.
 */
package org.tianyun.cloud.ext.thread;

import org.tianyun.cloud.context.RequestContext;
import org.tianyun.cloud.context.RequestContext.RequestContextHolder;
import org.tianyun.cloud.context.ServletContext.ServletContextHolder;

import java.util.Map;
import java.util.concurrent.Callable;

/**
 * 线程上下文
 *
 * @auther ebert_chan
 */
public class ThreadContext {

    /**
     * 包装Callable，使之支持MDC
     *
     * @param task
     * @return
     * @auther ebert_chan
     */
    public static <T> Callable<T> wrap(Callable<T> task) {
        RequestContext requestContext = RequestContextHolder.get();
        Map<String, Object> contextMap = ServletContextHolder.getCopyOfContextMap();
        return () -> {

            if (requestContext == null) {
                RequestContextHolder.clear();
            } else {
                RequestContextHolder.set(requestContext);
            }

            if (contextMap == null) {
                ServletContextHolder.clear();
            } else {
                ServletContextHolder.setContextMap(contextMap);
            }

            try {
                return task.call();
            } finally {
                RequestContextHolder.clear();
                ServletContextHolder.clear();
            }
        };
    }

    /**
     * 包装Runnable，使之支持MDC
     *
     * @param task
     * @return
     * @auther ebert_chan
     */
    public static Runnable wrap(Runnable task) {
        RequestContext requestContext = RequestContextHolder.get();
        Map<String, Object> contextMap = ServletContextHolder.getCopyOfContextMap();
        return () -> {

            if (requestContext == null) {
                RequestContextHolder.clear();
            } else {
                RequestContextHolder.set(requestContext);
            }

            if (contextMap == null) {
                ServletContextHolder.clear();
            } else {
                ServletContextHolder.setContextMap(contextMap);
            }
            try {
                task.run();
            } finally {
                RequestContextHolder.clear();
                ServletContextHolder.clear();
            }
        };
    }

}
